package com.example.demo.boot.util;

import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.ThreadPoolExecutor;

public class ThreadPoolUtils {

    public static ThreadPoolTaskExecutor createThreadPool(int poolSize, int queueCapacity, String threadNamePrefix) {
        ThreadPoolTaskExecutor threadPool = new ThreadPoolTaskExecutor();
        threadPool.setCorePoolSize(poolSize);
        threadPool.setMaxPoolSize(poolSize);
        threadPool.setQueueCapacity(queueCapacity);
        threadPool.setKeepAliveSeconds(60);
        threadPool.setAllowCoreThreadTimeOut(true);
        threadPool.setThreadNamePrefix(threadNamePrefix);
        threadPool.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        threadPool.setWaitForTasksToCompleteOnShutdown(true);
        threadPool.setAwaitTerminationSeconds(60);
        return threadPool;
    }

    public static ThreadPoolTaskExecutor createBlockingThreadPool(int poolSize, int queueCapacity, String threadNamePrefix) {
        ThreadPoolTaskExecutor threadPool = new ThreadPoolTaskExecutor();
        threadPool.setCorePoolSize(poolSize);
        threadPool.setMaxPoolSize(poolSize);
        threadPool.setQueueCapacity(queueCapacity);
        threadPool.setKeepAliveSeconds(60);
        threadPool.setAllowCoreThreadTimeOut(true);
        threadPool.setThreadNamePrefix(threadNamePrefix);
        threadPool.setRejectedExecutionHandler((r, executor) -> {
            if (!executor.isShutdown()) {
                try {
                    executor.getQueue().put(r);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            }
        });
        threadPool.setWaitForTasksToCompleteOnShutdown(true);
        threadPool.setAwaitTerminationSeconds(60);
        return threadPool;
    }


}